import torch
import torchvision.datasets as dataset
import torchvision.transforms as transforms
import torch.utils.data as datautils

# from cnn import CNN

# 指定路径已经有了，download不会下载
test_data = dataset.MNIST(
    "mnist",
    train=False,
    transform=transforms.ToTensor(),
    download=True,
)

test_loader = datautils.DataLoader(
    test_data,
    batch_size=64,
    shuffle=True
)

cnn = torch.load("model/mnist_cls.pkl")
cnn = cnn.cuda()

# loss
# eval/test
loss_test = 0
accuracy = 0

# pip install opencv-python -i http://mirrors.aliyun.com/pypi/simple/   --trusted-host mirrors.aliyun.com
# opencv绘图
import cv2

show_image = True
# 每次取出batchsize张图片
for i, (images, labels) in enumerate(test_loader):
    images = images.cuda()
    labels = labels.cuda()
    outputs = cnn(images)
    _, pred = outputs.max(1)
    accuracy += (pred == labels).sum().item()

    if show_image:
        # 绘图
        images = images.cpu().numpy()
        labels = labels.cpu().numpy()
        pred = pred.cpu().numpy()

        # images.shape:  batchsize * 1 * 28 *28
        for idx in range(images.shape[0]):
            im_data = images[idx]
            im_label = labels[idx]
            im_pred = pred[idx]
            # numpy表示图像时，通道放最后
            im_data = im_data.transpose(1,2,0)

            print("label: ", im_label, "; pred: ", im_pred)
            cv2.imshow("imdata", im_data)
            cv2.waitKey(0)


accuracy = accuracy/len(test_data)
print(accuracy)
